#!/usr/bin/env python
#-*- coding:utf-8 -*-
#Author='He Rensheng'
#Email='hrs323@126.com'

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import pandas as pd
import cv2, os, time, shutil
import numpy as np


def get_pic_after_eval(imagesdir, imagesetfile, csvfile, eval_ovthresh):
    # get filename list from imagesetfile
    with open(imagesetfile, 'r') as f:
        lines = f.readlines()
    imagenames = [x.strip() for x in lines]

    # load the csvfile to DataFrame
    dets = pd.read_csv(csvfile, dtype={'ImgName': object, 'ClsName': object})

    # 设置画好框的存放路径
    current_time = time.strftime('%Y-%m-%d_%H-%M-%S', time.localtime(time.time()))
    test_result_savepath = os.path.join("data/test_result", current_time + '_after-eval_ovthresh_' + str(eval_ovthresh))
    if not os.path.isdir(test_result_savepath):
        os.makedirs(test_result_savepath)

    # 遍历图片集
    for imagename in imagenames:
        # read image to ndarray
        im_file = os.path.join(imagesdir, imagename + '.jpg')
        im = cv2.imread(im_file)
        # print(im.shape)

        # ###### ----Newly added----Set the Savepath------ ############
        im = im[:, :, (2, 1, 0)]
        # 函数返回一个figure图像和一个子图ax的array列表; figsize定义画布大小, 与原图保持一致
        fig, ax = plt.subplots(figsize=(im.shape[1] * 0.01, im.shape[0] * 0.01))
        ax.imshow(im)
        # 使得图片充满整个画布,从而使得输出的图片与原图一样
        plt.subplots_adjust(top=1, bottom=0, left=0, right=1, hspace=0, wspace=0)

        # get the box info for this imagename from dets
        image_det = dets[dets['ImgName'].isin([imagename])]
        # print(image_det)

        for i in range(len(image_det)):
            # get the box info in the image_dets
            bbox = np.array([image_det.iloc[i]['xmin'], image_det.iloc[i]['ymin'],
                             image_det.iloc[i]['xmax'], image_det.iloc[i]['ymax']]).astype('float64')
            clsname = image_det.iloc[i]['ClsName']
            score = image_det.iloc[i]['score']
            state = image_det.iloc[i]['state']

            # 画框
            if state == 1:
                # TP
                ax.add_patch(
                    plt.Rectangle((bbox[0], bbox[1]),
                                  bbox[2] - bbox[0],
                                  bbox[3] - bbox[1], fill=False,
                                  edgecolor='green', linewidth=3.5)
                )
            elif state == 0:
                # FP
                ax.add_patch(
                    plt.Rectangle((bbox[0], bbox[1]),
                                  bbox[2] - bbox[0],
                                  bbox[3] - bbox[1], fill=False,
                                  edgecolor='red', linewidth=3.5)
                )
            elif state == -1:
                # FN
                ax.add_patch(
                    plt.Rectangle((bbox[0], bbox[1]),
                                  bbox[2] - bbox[0],
                                  bbox[3] - bbox[1], fill=False,
                                  edgecolor='magenta', linewidth=3.5)
                )
            else:
                print('image-{} does not has box info in csv'.format(imagename))

            ax.text(bbox[0], bbox[1] - 2, '{:s} {:.3f}'.format(clsname[0], score),
                    bbox=dict(facecolor='blue', alpha=0.3),
                    fontsize=12, color='white')

        plt.axis('off')  # 不显示坐标尺寸
        # plt.tight_layout()  # 紧凑显示图片，居中显示
        plt.draw()
        plt.savefig(os.path.join(test_result_savepath, imagename + '.jpg'))
        plt.close()

    # 把test阶段生成的test.csv文件mv到test_savepath下, 以备不时之需
    shutil.move(csvfile, test_result_savepath)


if __name__ == '__main__':
    imagesdir = 'data/VOCdevkit2007/VOC2007/JPEGImages'
    imagesetfile = 'data/VOCdevkit2007/VOC2007/ImageSets/Main/test.txt'
    csvfile = 'data/test_result/after-eval.csv'
    get_pic_after_eval(imagesdir, imagesetfile, csvfile)
    print('')
    print('-----The boxes has been done after evaluations!-----')
